package 题目集.二分;

import org.junit.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * https://leetcode.cn/problems/minimum-operations-to-make-all-array-elements-equal/description/
 * 二分+前缀和
 */
public class demo07_使数组元素全部相等的最少操作次数 {

    /**
     * 暴力思路：嵌套遍历q和n两个数组，计算n数组中每个数和q[i]的差值和。
     * 时间复杂度：n^2
     * 二分加前缀和思路：先对n数组进行排序，计算排序后的前缀和，然后二分查找n数组中小于和大于q[i]的一个数
     * 通过前缀和计算 左部分和右部分 的和。计算与q[i]*他们长度的差
     */
    int[] nums;
    long[] sum;

    public List<Long> minOperations(int[] nums, int[] queries) {
        Arrays.sort(nums);
        this.nums = nums;
        sum = new long[nums.length + 1];
        for (int i = 1; i < sum.length; i++) {
            sum[i] = sum[i - 1] + nums[i - 1];
        }
        List<Long> res = new ArrayList<>(queries.length);
        for (int i = 0; i < queries.length; i++) {
            int l = 0, r = nums.length - 1;
            while (l <= r) {
                int m = l + r >> 1;
                if (nums[m] > queries[i]) {
                    r = m - 1;
                } else {
                    l = m + 1;
                }
            }
            int left = r;
            int right = l;
            long rsum = sum[nums.length] - sum[right];
            long lsum = sum[left + 1] - sum[0];
            long lcnt = (long) (left + 1) * queries[i] - lsum;
            long rcnt = rsum - (long) (nums.length - right) * queries[i];
            res.add(lcnt + rcnt);
        }
        return res;
    }


    @Test
    public void test() {
        int[] nums = {1, 1, 2, 2, 3, 3};
        int[] q = {2};
        List<Long> longs = minOperations2(nums, q);
        System.out.println(longs);
    }

    public List<Long> minOperations2(int[] nums, int[] queries) {
        List<Long> ans = new ArrayList<>();
        // 排序nums
        Arrays.sort(nums);
        // 前缀和
        long[] sum = new long[nums.length + 1];
        sum[0] = 0;
        for (int i = 0; i < nums.length; i++) {
            sum[i + 1] = sum[i] + nums[i];
        }
        // 遍历
        for (int i = 0; i < queries.length; i++) {
            // 二分
            int left = -1;
            int right = nums.length;
            while (left + 1 != right) {
                int mid = left + (right - left) / 2;
                if (nums[mid] <= queries[i]) {
                    left = mid;
                } else {
                    right = mid;
                }
            }
            long left_area = left_area = (long) queries[i] * (left + 1) - (sum[left + 1] - sum[0]);
            ;
            long right_area = right_area = sum[nums.length] - sum[right] - (long) queries[i] * (nums.length - right);
            ans.add(left_area + right_area);
        }
        return ans;
    }
}
